use std::fs::{self, File};
use std::io::Write;
use std::path::{Path, PathBuf};
use text_splitter::{MarkdownSplitter, TextSplitter, CodeSplitter};

/// 将Markdown文本分割成多个块，每个块的token数量不超过指定的最大值
///
/// # 参数
///
/// * `text` - 要分割的Markdown文本
/// * `max_tokens` - 每个块的最大token数量
///
/// # 返回值
///
/// 返回分割后的Markdown块列表
pub fn split_markdown(text: &str, max_tokens: usize) -> Vec<String> {
    // 使用字符数量作为分割依据
    let splitter = TextSplitter::new(max_tokens);

    // 分割文本并收集结果
    splitter.chunks(text).map(|s| s.to_string()).collect()
}

/// 将Markdown文本分割成多个块，每个块的token数量不超过指定的最大值
/// 使用MarkdownSplitter，可以更好地保持Markdown的语义结构
///
/// # 参数
///
/// * `text` - 要分割的Markdown文本
/// * `max_tokens` - 每个块的最大token数量
///
/// # 返回值
///
/// 返回分割后的Markdown块列表
pub fn split_markdown_semantic(text: &str, max_tokens: usize) -> Vec<String> {
    // 使用MarkdownSplitter进行语义分割
    let splitter = MarkdownSplitter::new(max_tokens);

    // 分割文本并收集结果
    splitter.chunks(text).map(|s| s.to_string()).collect()
}

/// 将分割后的Markdown块保存到指定目录
///
/// # 参数
///
/// * `chunks` - 分割后的Markdown块列表
/// * `output_dir` - 输出目录路径
/// * `base_filename` - 基础文件名
///
/// # 返回值
///
/// 返回保存的文件路径列表
pub fn save_markdown_chunks(chunks: &[String], output_dir: &Path, base_filename: &str) -> Vec<PathBuf> {
    // 确保输出目录存在
    fs::create_dir_all(output_dir).expect("Failed to create output directory");

    let mut file_paths = Vec::new();

    // 保存每个块到单独的文件
    for (i, chunk) in chunks.iter().enumerate() {
        let filename = format!("{}_chunk_{}.md", base_filename, i + 1);
        let file_path = output_dir.join(&filename);

        let mut file = File::create(&file_path).expect("Failed to create file");
        file.write_all(chunk.as_bytes()).expect("Failed to write to file");

        file_paths.push(file_path);
    }

    file_paths
}

/// 将多个Markdown块合并成一个完整的文档
///
/// # 参数
///
/// * `chunks` - 要合并的Markdown块列表
///
/// # 返回值
///
/// 返回合并后的Markdown文档
pub fn combine_markdown_chunks(chunks: &[String]) -> String {
    chunks.join("\n")
}

/// 将HTML文本分割成多个块，每个块的token数量不超过指定的最大值
/// 使用tree-sitter-html进行语法分析，保持HTML结构的完整性
///
/// # 参数
///
/// * `html` - 要分割的HTML文本
/// * `max_tokens` - 每个块的最大token数量
///
/// # 返回值
///
/// 返回分割后的HTML块列表
pub fn split_html_semantic(html: &str, max_tokens: usize) -> Vec<String> {
    // 使用CodeSplitter配合tree-sitter-html进行语义分割
    let splitter = CodeSplitter::new(
        tree_sitter_html::LANGUAGE,
        max_tokens
    ).expect("Invalid tree-sitter HTML language");

    // 分割文本并收集结果
    splitter.chunks(html).map(|s| s.to_string()).collect()
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::path::Path;
    use common::utils::count_tokens;

    /// 测试基本的Markdown分割功能
    #[test]
    fn test_split_markdown_basic() {
        let text = "# 标题\n\n这是一段测试文本，用于测试Markdown分割功能。\n\n## 子标题\n\n更多的测试内容。";

        // 设置一个较小的token限制，确保文本会被分割
        let max_tokens = 10;
        let chunks = split_markdown(text, max_tokens);

        // 验证文本被分割成多个块
        assert!(chunks.len() > 1);

        // 验证每个块的token数量不超过最大值
        for chunk in &chunks {
            let token_count = count_tokens(chunk).unwrap_or(0);
            assert!(token_count <= max_tokens);
        }

        // 验证合并后的文本与原始文本相似（可能有些空白字符的差异）
        let combined = combine_markdown_chunks(&chunks);
        assert_eq!(combined.replace("\n\n", "\n").trim(), text.replace("\n\n", "\n").trim());
    }

    /// 测试语义化的Markdown分割功能
    #[test]
    fn test_split_markdown_semantic() {
        let text = "# 标题\n\n这是一段测试文本，用于测试Markdown分割功能。\n\n## 子标题\n\n更多的测试内容。";

        // 设置一个较小的token限制，确保文本会被分割
        let max_tokens = 10;
        let chunks = split_markdown_semantic(text, max_tokens);

        // 验证文本被分割成多个块
        assert!(chunks.len() > 1);

        // 验证每个块的token数量不超过最大值
        for chunk in &chunks {
            let token_count = count_tokens(chunk).unwrap_or(0);
            assert!(token_count <= max_tokens);
        }
    }

    /// 测试保存Markdown块到文件
    #[test]
    fn test_save_markdown_chunks() {
        let chunks = vec![
            "# 第一块".to_string(),
            "# 第二块".to_string(),
            "# 第三块".to_string(),
        ];

        let output_dir = Path::new("../../output/chunks/test");
        let file_paths = save_markdown_chunks(&chunks, output_dir, "test");

        // 验证文件数量与块数量相同
        assert_eq!(file_paths.len(), chunks.len());

        // 验证文件内容
        for (i, path) in file_paths.iter().enumerate() {
            let content = fs::read_to_string(path).expect("Failed to read file");
            assert_eq!(content, chunks[i]);
        }

        // 清理测试文件
        for path in file_paths {
            fs::remove_file(path).ok();
        }
        fs::remove_dir(output_dir).ok();
    }

    /// 测试使用真实的Markdown文档进行分割和重组
    #[test]
    fn test_split_and_combine_real_markdown() {
        // 读取真实的Markdown文档
        let markdown_path = Path::new("/data/projects/yslib/output/csug_gemini/zh/libraries.md");
        let text = fs::read_to_string(markdown_path).expect("Failed to read markdown file");

        // 设置最大token长度为5000
        let max_tokens = 5000;

        // 使用语义分割方法
        let chunks = split_markdown_semantic(&text, max_tokens);

        // 验证分割结果
        assert!(!chunks.is_empty());

        // 验证每个块的token数量不超过最大值
        for (i, chunk) in chunks.iter().enumerate() {
            let token_count = count_tokens(chunk).unwrap_or(0);
            println!("Chunk {}: {} tokens", i + 1, token_count);
            assert!(token_count <= max_tokens);
        }

        // 保存分割后的块到文件
        let output_dir = Path::new("/data/projects/yslib/output/chunks/libraries");
        let file_paths = save_markdown_chunks(&chunks, output_dir, "libraries");

        // 验证文件已保存
        assert_eq!(file_paths.len(), chunks.len());

        // 合并块并验证结果
        let combined = combine_markdown_chunks(&chunks);

        // 保存合并后的文档以便手动检查
        let combined_path = output_dir.join("libraries_combined.md");
        fs::write(&combined_path, &combined).expect("Failed to write combined file");

        // 验证合并后的文档与原始文档的token数量相近
        let original_tokens = count_tokens(&text).unwrap_or(0);
        let combined_tokens = count_tokens(&combined).unwrap_or(0);

        println!("Original document: {} tokens", original_tokens);
        println!("Combined document: {} tokens", combined_tokens);

        // 允许有少量差异（由于空白字符处理等原因）
        let difference = if original_tokens > combined_tokens {
            original_tokens - combined_tokens
        } else {
            combined_tokens - original_tokens
        };

        // 差异不应超过原始文档token数量的5%
        let max_allowed_difference = original_tokens as f64 * 0.05;
        assert!(difference as f64 <= max_allowed_difference,
                "Token count difference too large: {}", difference);
    }

    #[test]
    fn test_split_html_semantic() {
        let html = r#"
            <!DOCTYPE html>
            <html>
            <head>
                <title>Test Page</title>
            </head>
            <body>
                <h1>Main Title</h1>
                <div class="content">
                    <p>First paragraph with some content.</p>
                    <p>Second paragraph with different content.</p>
                </div>
            </body>
            </html>
        "#;

        // 设置一个较小的token限制，确保HTML会被分割
        let max_tokens = 50;
        let chunks = split_html_semantic(html, max_tokens);

        // 验证文本被分割成多个块
        assert!(chunks.len() > 1);

        // 验证每个块的token数量不超过最大值
        for chunk in &chunks {
            println!("Chunk: {}", chunk);
            let token_count = count_tokens(chunk).unwrap_or(0);
            assert!(token_count <= max_tokens);
        }

        // 验证每个块都是有效的HTML片段
        for chunk in &chunks {
            assert!(chunk.contains("<") && chunk.contains(">"));
        }
    }

    /// 测试使用真实的HTML文档进行分割，并将分割后的块保存到文件
    #[test]
    fn test_split_real_html_semantic() {
        // 读取真实的HTML文档
        let html_path = Path::new("/data/projects/yslib/resources/csug/foreign.html");
        let html = fs::read_to_string(html_path).expect("Failed to read HTML file");

        // 设置最大token长度为10000
        let max_tokens = 10000;

        // 使用语义分割方法
        let chunks = split_html_semantic(&html, max_tokens);

        // 验证分割结果
        assert!(!chunks.is_empty());
        println!("HTML split into {} chunks", chunks.len());

        // 验证每个块的token数量不超过最大值
        for (i, chunk) in chunks.iter().enumerate() {
            let token_count = count_tokens(chunk).unwrap_or(0);
            println!("Chunk {}: {} tokens", i + 1, token_count);
            assert!(token_count <= max_tokens, "Chunk {} exceeds max tokens: {}", i + 1, token_count);
        }

        // 保存分割后的块到文件
        let output_dir = Path::new("/data/projects/yslib/output/chunks/foreign_html");
        let file_paths = save_markdown_chunks(&chunks, output_dir, "foreign_html");

        // 验证文件已保存
        assert_eq!(file_paths.len(), chunks.len());
        println!("Saved {} HTML chunks to {}", chunks.len(), output_dir.display());

        // 计算原始HTML的token数量
        let original_tokens = count_tokens(&html).unwrap_or(0);
        println!("Original HTML document: {} tokens", original_tokens);

        // 计算所有块的总token数量
        let total_chunk_tokens: usize = chunks.iter()
            .map(|chunk| count_tokens(chunk).unwrap_or(0))
            .sum();
        println!("Total tokens in all chunks: {}", total_chunk_tokens);

        // 计算分割前后的token数量差异
        let difference = if original_tokens > total_chunk_tokens {
            original_tokens - total_chunk_tokens
        } else {
            total_chunk_tokens - original_tokens
        };
        println!("Token count difference: {} ({:.2}%)",
                 difference,
                 (difference as f64 / original_tokens as f64) * 100.0);
    }
}
